import json
import requests
import urllib.parse as urlparse
import time
import hashlib
import hmac
import sys

from coin.exchange.base.kr_rest.auth_base import AuthKey
from coin2.api_client.common import (SyncQueryer)
from enum import Enum

class Side(Enum):
  BUY = "buy"
  SELL = "sell"

class Type(Enum):
  LIMIT = "limit"
  MARKET = "market"
  LIMIT_MAKER = "limit_maker"
  IOC = "ioc"

class OrderMode(Enum):
  SPOT = "spot"
  ISO_MARGIN = "iso_margin"
  ALL = "all"

class OrderStatus(Enum):
  PENDING_ORDER = 4                   # Order success, Pending for fulfilment
  PARTIALLY_FILLED = 5                # Partially filled
  FULLY_FILLED = 6                    # Fully filled
  CANCELED = 8                        # Canceled
  OPEN_ORDER = 9                      # Outstanding (4 and 5)
  COMPELETE_ORDER = 10                # 6 and 8 and 11
  PARTIALLY_FILLED_AND_CANCELED = 11  # Partially filled and canceled

class BitmartAuth(requests.auth.AuthBase):
  def __init__(self, key):
    self._key = AuthKey.from_file(key)

  def _gen_signature(self, timestamp: str, memo: str, query_string: str):
    before_sign = (timestamp + "#" + memo + "#" + query_string).encode()

    return hmac.new(self._key.secret_key_bytes, before_sign, digestmod=hashlib.sha256).hexdigest()

  def __call__(self, req: requests.PreparedRequest):
    if req.method == "GET":
      query_string = urlparse.urlparse(req.url).query
    elif req.method == "POST":
      if isinstance(req.body, bytes):
        query_string = req.body.decode()
      else:
        query_string = req.body

    timestamp = str(int(time.time() * 1000))
    sign = self._gen_signature(timestamp, self._key.key_name, query_string)

    req.headers.update({"X-BM-KEY": self._key.access_key})
    req.headers.update({"X-BM-SIGN": sign})
    req.headers.update({"X-BM-TIMESTAMP": timestamp})
    req.headers.update({"Content-Type": "application/json"})

    req.prepare_headers(req.headers)

    # print(req.headers)
    # print(req.body)
    # print(req.url)

    return req

class BitmartPrivateClient():
  def __init__(self, key_path: str, ip=None):
    self._auth = BitmartAuth(key_path)
    self._url = "https://api-cloud.bitmart.com"

    self._query = SyncQueryer(api_host=self._url, auth=self._auth, ip=ip)

  def get_orders(
    self,
    symbol: str,
    status: OrderStatus,
    order_mode: OrderMode = None,
    limit: int = None,
    start_time: int = None,
    end_time: int = None,
  ):
    assert symbol, "symbol must be not None and not empty!"
    assert status, "status must be not None!"

    method = "GET"
    params = {
      'symbol': symbol,
      'order_mode': order_mode and order_mode.value,
      'N': limit,
      'start_time': start_time,
      'end_time': end_time,
      'status': status.value,
    }
    path = "/spot/v3/orders"

    response = self._query.query(method=method, path=path, params=params)
    return response.json()

  def get_trades(
    self,
    symbol: str,
    order_mode: OrderMode = None,
    order_id: str = None,
    limit: int = None,
    start_time: int = None,
    end_time: int = None,
  ):
    assert symbol, "symbol must be not None and not empty!"

    method = "GET"
    params = {
      'symbol': symbol,
      'order_mode': order_mode and order_mode.value,
      'order_id': order_id,
      'N': limit,
      'start_time': start_time,
      'end_time': end_time,
    }
    path = "/spot/v2/trades"

    response = self._query.query(method=method, path=path, params=params)
    return response.json()

  def cancel_order(
    self,
    symbol: str,
    order_id: str = None,
    client_order_id: str = None
  ):
    assert symbol, "symbol must be not None and not empty!"
    one_and_only_one = (order_id == None and client_order_id != None) or (order_id != None and client_order_id == None)
    assert one_and_only_one, "order_id and client_order_id have one and only one!"

    method = "POST"
    json = {
      'symbol': symbol,
      'order_id': order_id,
      'client_order_id': client_order_id,
    }
    path = "/spot/v3/cancel_order"

    response = self._query.query(method=method, path=path, json=json)
    return response.json()

  def submit_order(
    self,
    symbol: str,
    side: Side,
    type: Type,
    client_order_id: str = None,
    size: str = None,
    price: str = None,
    notional: str = None,
  ):
    assert symbol, "symbol must be not None and not empty!"
    assert side, "side must be not None!"
    assert type, "type must be not None!"

    method = "POST"
    json = {
      'symbol': symbol,
      'side': side.value,
      'type': type.value,
      'client_order_id': client_order_id,
      'size': size,
      'price': price,
      'notional': notional,
    }
    path = "/spot/v2/submit_order"

    response = self._query.query(method=method, path=path, json=json)
    return response.json()

  def get_account_balance(self, currency: str = None):
    method = "GET"
    params = {
      'currency': currency,
    }
    path = "/account/v1/wallet"

    response = self._query.query(method=method, path=path, params=params)
    return response.json()

# api doc url
# https://developer-pro.bitmart.com/en/spot/#introduction-2

if __name__ == "__main__":
  key_path = sys.argv[1]
  
  bitmart = BitmartPrivateClient(key_path)
  info = bitmart.get_orders("BTC_USDT", OrderStatus.OPEN_ORDER)
  # info = bitmart.get_trades("BTC_USDT")
  # info = bitmart.cancel_order("BTC_USDT", client_order_id="test1")
  # info = bitmart.submit_order("BTC_USDT", Side.SELL, Type.MARKET, "20230207test2", size="0.00029")
  # info = bitmart.get_account_balance()
  print(json.dumps(info, indent=2))
